/**
 * Copyright 2009 Erick Armbrust 
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0 
 *   
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License. 
 */

package tapioca.util;

import java.util.Collection;
import java.util.List;
import java.util.Map;


import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyRange;
import com.google.appengine.api.datastore.PreparedQuery;
import com.google.appengine.api.datastore.Query;
import com.google.appengine.api.datastore.Transaction;

public class DelegateDatastoreService implements DatastoreService {

  private final DatastoreService ds;
  private final DatastoreStats stats;
  
  public DelegateDatastoreService(DatastoreService ds) {
    this.ds = ds;
    stats = new DatastoreStats();
  }  
  
  public DatastoreStats getStats() {
    return stats;
  }
  
  private int getSizeOfIterable(Iterable<?> iterable) {
    if (iterable.getClass().isAssignableFrom(Collection.class)) {
      return Collection.class.cast(iterable).size();
    }
    
    int size = 0;
    for (@SuppressWarnings("unused") Object o : iterable) {
      size++;
    }
    return size;
  }
  
  public KeyRange allocateIds(Key parent, String kind, long num) {
    stats.addIdAllocations(1);
    return ds.allocateIds(parent, kind, num);
  }

  public KeyRange allocateIds(String kind, long num) {
    stats.addIdAllocations(1);
    return ds.allocateIds(kind, num);
  }

  public Transaction beginTransaction() {
    return ds.beginTransaction();
  }

  public void delete(Iterable<Key> keys) {
    stats.addDeleteRequests(1);
    stats.addDeleteEntities(getSizeOfIterable(keys));
    ds.delete(keys);
  }

  public void delete(Key... keys) {
    stats.addDeleteRequests(1);
    stats.addDeleteEntities(keys.length);
    ds.delete(keys);
  }

  public void delete(Transaction txn, Iterable<Key> keys) {
    stats.addDeleteRequests(1);
    stats.addDeleteEntities(getSizeOfIterable(keys));
    ds.delete(txn, keys);
  }

  public void delete(Transaction txn, Key... keys) {
    stats.addDeleteRequests(1);
    stats.addDeleteEntities(keys.length);
    ds.delete(txn, keys);
  }

  public Map<Key, Entity> get(Iterable<Key> keys) {
    stats.addGetRequests(1);
    stats.addGetEntities(getSizeOfIterable(keys));
    return ds.get(keys);
  }

  public Entity get(Key key) throws EntityNotFoundException {
    stats.addGetRequests(1);
    stats.addGetEntities(1);
    return ds.get(key);
  }

  public Map<Key, Entity> get(Transaction txn, Iterable<Key> keys) {
    stats.addGetRequests(1);
    stats.addGetEntities(getSizeOfIterable(keys));
    return ds.get(txn, keys);
  }

  public Entity get(Transaction txn, Key key) throws EntityNotFoundException {
    stats.addGetRequests(1);
    stats.addGetEntities(1);
    return ds.get(txn, key);
  }

  public Collection<Transaction> getActiveTransactions() {
    return ds.getActiveTransactions();
  }

  public Transaction getCurrentTransaction() {
    return ds.getCurrentTransaction();
  }

  public Transaction getCurrentTransaction(Transaction txn) {
    return ds.getCurrentTransaction(txn);
  }

  public PreparedQuery prepare(Query query) {
    return ds.prepare(query);
  }

  public PreparedQuery prepare(Transaction txn, Query query) {
    return ds.prepare(txn, query);
  }

  public Key put(Entity entity) {
    stats.addPutRequests(1);
    stats.addPutEntities(1);
    return ds.put(entity);
  }

  public List<Key> put(Iterable<Entity> entities) {
    stats.addPutRequests(1);
    stats.addPutEntities(getSizeOfIterable(entities));
    return ds.put(entities);
  }

  public Key put(Transaction txn, Entity entity) {
    stats.addPutRequests(1);
    stats.addPutEntities(1);
    return ds.put(txn, entity);
  }

  public List<Key> put(Transaction txn, Iterable<Entity> entities) {
    stats.addPutRequests(1);
    stats.addPutEntities(getSizeOfIterable(entities));
    return ds.put(txn, entities);
  }
  
}
